#!/usr/bin/env python

import psycopg2
import psycopg2.extras
import sys

class DBStats(object):

    def get_artist_count(self):
        s = """SELECT count(artist_id) 
                 FROM artist"""
        cur.execute(s)
        return cur.fetchone()[0]

    def get_tracks_count(self):
        s = """SELECT count(track_id)
                 FROM tracks"""

        cur.execute(s)
        return cur.fetchone()[0]

    def get_album_count(self):
        s = """SELECT count(album_id)
                 FROM album"""

        cur.execute(s)
        return cur.fetchone()[0]

    def get_genre_count(self):
        s = """SELECT count(name)
                 FROM genre"""

        cur.execute(s)
        return cur.fetchone()[0]

    def get_album_track_count(self, album_id):
        s = """SELECT count(*)
                 FROM tracks
                WHERE album_id %s"""

        cur.execute(s, (album_id,))
        return cur.fetchone

    def get_filepath_count(self):
        s = """SELECT count(path_id)
                 FROM path"""

        cur.execute(s)
        return cur.fetchone()[0]

    def stats_report(self):
        print "    Artists => %s " % self.get_artist_count()
        print "     Albums => %s " % self.get_album_count()
        print "     Tracks => %s " % self.get_tracks_count()
        print "     Genres => %s " % self.get_genre_count()
        print "Directories => %s " % self.get_filepath_count()

class DBSearch(object):
    
    def _escape_query(self, query):
        """
        builds the appropriate escape sequences to pass into the LIKE 
        clause
        """ 
        out = ""
        for x in query:
            x.replace('=', '==').replace('%', '=%').replace('_', '=_')
            out = out + '%' + x 
        return out + '%'

    def get_album_tracks(self, album_id):
        """
        Returns a dict of artist_name, track_title, file_name
        """ 
        s = """SELECT art.name AS "artist_name" , trk.title AS "track_title", trk.file_name AS "file_name"
                 FROM tracks trk, artist art, album_map amp
                WHERE amp.track_id = trk.track_id
                  AND art.artist_id = trk.artist_id
                  AND amp.album_id = %s"""
        
        cur.execute(s, (album_id,)) 
        return cur.fetchall()

    def get_album_info(self, album_id):
        """
        Returns a dict of album_name, year, path_name, artist_name, genre_name
        """
        s = """SELECT alb.name as "album_name", alb.year, pth.name as
                      "path_name", art.name as "artist_name", gen.name as "genre_name"
                 FROM album alb, path pth, artist art, genre gen
                WHERE alb.path_id = pth.path_id
                  AND alb.artist_id = art.artist_id
                  AND alb.genre_id = gen.genre_id
                  AND alb.album_id = %s"""
       
        cur.execute(s, (album_id,))
        return cur.fetchone()

    def album_title(self, query):
        """
        Search releases based on album title (album.name)

        Accepts a search string
        Returns a list of album ids (album.album_id)
        """    
        s = """SELECT alb.album_id
                 FROM artist art, album alb, path pth
                WHERE alb.path_id = pth.path_id
                  AND art.artist_id = alb.artist_id
                  AND alb.name ILIKE %(like)s ESCAPE '='"""

        cur.execute(s, dict(like = self._escape_query(query) ))
        return cur.fetchall()

    def artist_name(self, query):
        """
        Search releases bassed on the artist name (artist.name)

        Accepts a search string 
        Returns a list of album ids (album.album_id)
        """
        
        s = """SELECT alb.album_id
                 FROM artist art, album alb
                WHERE art.artist_id = alb.artist_id
                  AND art.name ILIKE %(like)s ESCAPE '='"""

        cur.execute(s, dict(like = self._escape_query(query)))
        return cur.fetchall()

    def track_name(self, query):

        s = """SELECT alb.album_id
                 FROM album alb, tracks trk, album_map amap
                WHERE amap.track_id = trk.track_id
                  AND amap.album_id = alb.album_id
                  AND trk.name ILIKE %(like)s ESCAPE '='"""

    def unified(self, query):
        """
        Executes a %LIKE% query on album.name OR artist.name OR track.title.
        """
        s ="""SELECT distinct(alb.album_id), alb.year
                FROM album_map amap, artist art, album alb, tracks trk
               WHERE art.artist_id = trk.artist_id
                 AND amap.track_id = trk.track_id
                 AND amap.album_id = alb.album_id
                 AND (alb.name ILIKE %(like)s ESCAPE '='
                  OR art.name ILIKE %(like)s ESCAPE '='
                  OR trk.title ILIKE %(like)s ESCAPE '=') 
            ORDER BY alb.year"""

        cur.execute(s, dict(like = self._escape_query(query)))
        return cur.fetchall()

    def new_additions(self, n):
        """
        Returns latest N number of additions to the filesystem.
        path.create_date is obtained from the last modified timestamp
        """
        s = """SELECT alb.album_id
                FROM album alb, path pth
               WHERE alb.path_id = pth.path_id
                 AND pth.create_date >= (now() - interval '%s days')
            ORDER BY pth.create_date""" % n
        cur.execute(s)
        return cur.fetchall()

if __name__ == "__main__":
    from optparse import OptionParser
    import sys
    
    p = OptionParser()
    p.add_option("-q", "--query", action="store", dest="query",
                 help="String to query against")
    p.add_option("-s", "--stats", action="store_true", dest="stats",
                 help="Display database statistics")
    p.add_option("-n", "--new", action="store", dest="new",
                 help="Display new additions for N number of days")
    p.add_option("-t", "--terse", action="store_true", dest="terse",
                 help="Display terse formatted output")
    p.add_option("-m", "--method", action="store", dest="method",
                 help="Type of search to perform, defaults to all \
                        Valid options are : artist, track,\
                        album or all")

    p.set_defaults(stats = False, terse=False, method='all')
    (options, args) = p.parse_args()

    conn = psycopg2.connect("dbname=mp3 user=mp3 host=192.168.0.1 \
                             password=mp3")
    cur = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
   
    if options.stats: 
        dbs = DBStats()
        dbs.stats_report()
        sys.exit(0)
   
    if not options.query and not options.new:
        p.error("ERROR : You must pass a query string (-q)")

    search = DBSearch()
    
    if options.new:
        albums = search.new_additions(options.new)
    elif options.query:
        query = options.query.split(" ")
        
        if options.method == "artist":
            albums = search.artist_name(query)
        elif options.method == "track":
            albums = search.track_name(query)
        elif options.method == "album":
            albums = search.album_title(query)
        else:
            albums = search.unified(query)
    
        print "[%s search] %d albums found matching term(s)  %s\n" % \
        (options.method, len(albums), ', '.join('%s' % q for q in query))

    # loop through each album_id
    for a in albums:
        # retrieve data for specific ID.
        alb = search.get_album_info(a[0])
        print "%s - %s [%s / %s]" % (alb['artist_name'], alb['album_name'], \
                alb['genre_name'], alb['year'])
        tracks = search.get_album_tracks(a[0])
        # tracks are not displayed in terse mode..
        if not options.terse:
            print "%s" % alb['path_name']
            for t in tracks:
                print "\t %s - %s" % (t['artist_name'], t['track_title']) 
